import shutil
import os.path
import sys
from kuai.analysis import diffmol
from kuai.dff import dffopt, dffqmd
import kuai

from group import GROUPS 

DFF_BIN = r"D:\dff7\bin32w"
PPF_FILE = r"D:\DATA\DffForceField\T_amber_ver3\fit\group2\ff\outputAMBER.ppf"

DATA_FOLDER = r"D:\DATA\DffForceField\T_amber_ver3\fit\group2\intra\Data"

def name_handle(msdfile):
    if msdfile.split('\\') > 2:
        a = msdfile.split('\\')
        new_name = a[-2]+'_' + a[-1]
    return new_name

def scan_file(qmdfile, msdfile):
    try:
        mol1 = kuai.read_mol(msdfile)
        mol1.name = msdfile[len(DATA_FOLDER):]
        qmd = kuai.qmd.get_opt_conform(kuai.read_qmd(open(qmdfile)))
        qmd.update_mol(mol1)
        kuai.assign_charges(mol1, ff, eqt)
        mol2 = dffopt.optimize(mol1, ff, {"DFF_BIN":DFF_BIN})
        #these lines are added to output the msd file after opt
        msd_output = name_handle(msdfile)
        kuai.write_mol("D:\DATA\DffForceField\intra\intra\compare\opt\%s"%msd_output, mol2)
        # end of lines
        diffmol.compare_mol(mol1, mol2)
    except RuntimeError, error:
        sys.stderr.write("#Error@" + msdfile[len(DATA_FOLDER):] + ": " + str(error)+"\n")

def scan_data(folder):
    subs = os.listdir(folder)
    for i in subs:
        if i == "Later":
            continue
        fullname = os.path.join(folder, i)
        if os.path.isdir(fullname):
            scan_data(fullname)
        elif os.path.isfile(fullname) and i.endswith(".qmd"):
            msdfile = fullname[:-4] + ".msd"
            if os.path.isfile(msdfile):
                scan_file(fullname, msdfile)

def scan_data_groups(folder):
    subs = os.listdir(folder)
    for i in subs:
        if i in GROUPS:
            if i == "Later":
                continue
            fullname = os.path.join(folder, i)
            if os.path.isdir(fullname):
                scan_data(fullname)
            elif os.path.isfile(fullname) and i.endswith(".qmd"):
                msdfile = fullname[:-4] + ".msd"
                if os.path.isfile(msdfile):
                    scan_file(fullname, msdfile)


if __name__ == '__main__':
    shutil.copyfile(PPF_FILE, "./.dff.ppf")
    shutil.copyfile("%s.eqt"%PPF_FILE[:-4], "./.dff.eqt")
    ff, eqt = kuai.read_ff(PPF_FILE)
    #scan_data_groups(DATA_FOLDER)
